# File: fedentgate/core/dp_noise_injection.py
"""Differential Privacy Noise Injection Module"""

import torch
from opacus.accountants import RDPAccountant

class DPNoiseGenerator:
    def __init__(self, sigma_low=0.5, sigma_high=2.0, clip_norm=1.0):
        """
        Initialize dynamic noise injection parameters
        Args:
            sigma_low: Noise scale for activated clients (σ_low)
            sigma_high: Noise scale for suppressed clients (σ_high)
            clip_norm: Gradient clipping norm (C)
        """
        self.sigma_low = sigma_low
        self.sigma_high = sigma_high
        self.clip_norm = clip_norm
        self.accountant = RDPAccountant()

    def clip_gradients(self, gradients):
        """Apply gradient clipping with L2 norm"""
        grad_norm = gradients.norm(2).item()
        scale = min(1.0, self.clip_norm / (grad_norm + 1e-6))
        return gradients * scale

    def add_noise(self, gradients, gating_state):
        """
        Add adaptive Gaussian noise based on gating state
        Args:
            gradients: Client's raw gradients
            gating_state: Binary gating decision (0/1)
        Returns:
            noisy_gradients: Noise-masked gradients
            sigma: Actual noise scale applied
        """
        clipped_grads = self.clip_gradients(gradients)
        sigma = self.sigma_low if gating_state == 1 else self.sigma_high
        noise = torch.randn_like(clipped_grads) * sigma
        return clipped_grads + noise, sigma

    def track_privacy_cost(self, sigma, sampling_rate):
        """
        Track privacy expenditure using RDP accountant
        Args:
            sigma: Noise scale used
            sampling_rate: Client sampling probability
        """
        self.accountant.step(noise_multiplier=sigma, sample_rate=sampling_rate)

    def get_privacy_spent(self, delta=1e-5):
        """Convert RDP to (ε, δ)-DP"""
        return self.accountant.get_privacy_spent(delta=delta)
